/*
//
//              INTEL CORPORATION PROPRIETARY INFORMATION
//  This software is supplied under the terms of a license  agreement or
//  nondisclosure agreement with Intel Corporation and may not be copied
//  or disclosed except in  accordance  with the terms of that agreement.
//    Copyright (c) 2008 Intel Corporation. All Rights Reserved.
//
//
*/

#include "umc_defs.h"
#if defined(UMC_ENABLE_AVS_SPLITTER)

#include "umc_avs_splitter_buffer.h"
#include "umc_automatic_mutex.h"
#include "umc_media_data_ex.h"
#include "vm_debug.h"

namespace UMC
{

enum
{
    DEFAULT_NUMBER_OF_SAMPLE    = 9,
    MIN_WRITE_PORTION           = 64 * 1024,
    EXTRA_MEMORY                = 1024,
    ALIGN_VALUE                 = 128
};

enum
{
    LAST_SLICE_START_CODE       = 0x0af,
    VIDEO_SEQUENCE_START_CODE   = 0x0b0,
    VIDEO_SEQUENCE_END_CODE     = 0x0b1,
    USER_DATA_START_CODE        = 0x0b2,
    I_PICTURE_START_CODE        = 0x0b3,
    EXTENSION_START_CODE        = 0x0b5,
    BP_PICTURE_START_CODE       = 0x0b6,
    VIDEO_EDIT_CODE             = 0x0b7,
    FIRST_SYSTEM_START_CODE     = 0x0b9
};

AVSSampleBuffer::AVSSampleBuffer(void)
{
    m_pbBuffer = NULL;
    m_nBufferSize = 0;
    m_nInputSize = 0;

    m_pbSample = NULL;
    m_pbSampleEnd = NULL;

    m_iNumberOfStartCodes = 0;

    m_bFoundSequenceHeader = false;

} // AVSSampleBuffer::AVSSampleBuffer(void)

AVSSampleBuffer::~AVSSampleBuffer(void)
{
    Close();

} // AVSSampleBuffer::~AVSSampleBuffer(void)

Status AVSSampleBuffer::Close(void)
{
    // release object's resources
    if (m_pbBuffer)
    {
        ippsFree(m_pbBuffer);
    }

    m_pbBuffer = NULL;
    m_nBufferSize = 0;
    m_nInputSize = 0;

    m_pbSample = NULL;
    m_pbSampleEnd = NULL;

    m_iNumberOfStartCodes = 0;

    m_bFoundSequenceHeader = false;

    // call the parent's method
    MediaBuffer::Close();

    return UMC_OK;

} // Status AVSSampleBuffer::Close(void)

Status AVSSampleBuffer::Init(MediaReceiverParams *init)
{
    MediaBufferParams *pParams = DynamicCast<MediaBufferParams> (init);
    Status umcRes;

    // check error(s)
    if (NULL == pParams)
        return UMC_ERR_NULL_PTR;

    // release the object before the initialization
    Close();

    // check parameter(s) & generate the default
    if ((0 == pParams->m_numberOfFrames) ||
        (0 == pParams->m_prefInputBufferSize) ||
        (0 == pParams->m_prefOutputBufferSize))
    {
        m_nInputSize = MIN_WRITE_PORTION;

        m_bufferParams.m_numberOfFrames = DEFAULT_NUMBER_OF_SAMPLE;
        m_bufferParams.m_prefInputBufferSize = MIN_WRITE_PORTION * 2;
        m_bufferParams.m_prefOutputBufferSize = MIN_WRITE_PORTION * 2;
        m_bufferParams.m_pMemoryAllocator = pParams->m_pMemoryAllocator;
    }
    else
    {
        m_bufferParams = *pParams;

        m_nInputSize = pParams->m_prefInputBufferSize;

        m_bufferParams.m_prefInputBufferSize *= 2;
    }

    // call the parent's method
    umcRes = MediaBuffer::Init(&m_bufferParams);
    if (UMC_OK != umcRes)
        return umcRes;

    //
    // do specific initializations here
    //

    // initialize start code array
    umcRes = m_startCodes.Init();
    if (UMC_OK != umcRes)
        return umcRes;

    // allocate buffer for the being collected sample
    m_nBufferSize = m_nInputSize * 2;
    m_pbBuffer = ippsMalloc_8u((int) (m_nBufferSize + EXTRA_MEMORY));
    if (NULL == m_pbBuffer)
        return UMC_ERR_ALLOC;
    m_pbSample = m_pbBuffer;
    m_pbSampleEnd = m_pbBuffer;

    // initialize the main buffer
    umcRes = m_mainBuffer.Init(&m_bufferParams);
    if (UMC_OK != umcRes)
        return UMC_OK;

    return UMC_OK;

} // Status AVSSampleBuffer::Init(MediaReceiverParams *pParams)

Status AVSSampleBuffer::LockInputBuffer(MediaData *pDst)
{
    size_t nFreeSize;
    Status umcRes;

    // check error(s)
    if (NULL == pDst)
        return UMC_ERR_NULL_PTR;
    if (NULL == m_pbBuffer)
        return UMC_ERR_NOT_INITIALIZED;

    // move collected sample
    // into the main buffer
    umcRes = MoveCollectedSample(UMC_OK);
    if (UMC_OK != umcRes)
        return umcRes;

    // compute free size in the temporal buffer
    nFreeSize = m_nBufferSize - (m_pbSampleEnd - m_pbSample);

    // we have filled the whole temporal buffer,
    // but a sample has not been collected.
    // we need to enlarge the temporal buffer
    if (m_nInputSize > nFreeSize)
    {
        size_t nSize = m_nBufferSize * 2;
        Ipp8u *pbTemp;

        // allocate a new buffer
        pbTemp = ippsMalloc_8u((int) (nSize + EXTRA_MEMORY));
        if (NULL == pbTemp)
            return UMC_ERR_ALLOC;

        // copy there existing data and exchange the pointers
        memcpy(pbTemp, m_pbSample, m_pbSampleEnd - m_pbSample);
        m_pbSampleEnd = pbTemp + (m_pbSampleEnd - m_pbSample);
        m_pbSample = pbTemp;
        ippsFree(m_pbBuffer);
        m_pbBuffer = pbTemp;
        m_nBufferSize = nSize;
    }

    // reserve space in the main buffer (maximum possible sample)
    {
        size_t nMaxSampleSize;

        nMaxSampleSize = (m_pbSampleEnd - m_pbSample) +
                         sizeof(Ipp32u) * (m_iNumberOfStartCodes * 2 + 2);

        umcRes = m_mainBuffer.LockInputBuffer(pDst, nMaxSampleSize);
        if (UMC_OK != umcRes)
            return umcRes;
    }

    // set the destination pointer
    pDst->SetBufferPointer(m_pbSampleEnd, m_nInputSize);

    return UMC_OK;

} // Status AVSSampleBuffer::LockInputBuffer(MediaData *pDst)

Status AVSSampleBuffer::UnLockInputBuffer(MediaData *pDst, Status StreamStatus)
{
    size_t nFreeSize;
    size_t nDataSize = 0;

    // check error(s)
    if (NULL == m_pbBuffer)
        return UMC_ERR_NOT_INITIALIZED;

    // get free size
    nFreeSize = m_nBufferSize - (m_pbSampleEnd - m_pbSample);

    if (pDst)
    {
        if (pDst->GetDataSize() > nFreeSize)
            return UMC_ERR_NOT_ENOUGH_BUFFER;

        nDataSize = pDst->GetDataSize();
    }

    // add 'video sequence end' code
    if ((UMC_OK != StreamStatus) &&
        (m_iNumberOfStartCodes) &&
        (VIDEO_SEQUENCE_END_CODE != m_startCodes[m_iNumberOfStartCodes - 1].m_startCode))
    {
        Ipp8u *pbTmp = m_pbSampleEnd + nDataSize;

        pbTmp[0] = 0;
        pbTmp[1] = 0;
        pbTmp[2] = 1;
        pbTmp[3] = VIDEO_SEQUENCE_END_CODE;

        nDataSize += sizeof(Ipp32u);
    }

    // process all start codes
    CollectStartCodes((Ipp32s) nDataSize);

    // this is the only case when we move collected sample
    // into the main buffer in this function
    if (UMC_OK != StreamStatus)
        MoveCollectedSample(StreamStatus);

    return UMC_OK;

} // Status AVSSampleBuffer::UnLockInputBuffer(MediaData *pDst, Status StreamStatus)

Status AVSSampleBuffer::LockOutputBuffer(MediaData *pSrc)
{
    MediaDataEx *pSrcEx;
    Status umcRes;
    Ipp32u *pCodes;

    // check error(s)
    if (NULL == pSrc)
        return UMC_ERR_NULL_PTR;
    if (NULL == m_pbBuffer)
        return UMC_ERR_NOT_INITIALIZED;

    // lock the space in the main buffer
    umcRes = m_mainBuffer.LockOutputBuffer(&m_latestSample);
    if (UMC_OK != umcRes)
        return umcRes;

    // set pointer to start codes
    pCodes = (Ipp32u *) (((Ipp8u *) m_latestSample.GetDataPointer()) + m_latestSample.GetDataSize());

    // set destination media data
    pSrc->SetBufferPointer((Ipp8u *) m_latestSample.GetBufferPointer(), pCodes[-1]);
    pSrc->SetDataSize(pCodes[-1]);

    // try to cast to the extended type
    pSrcEx = DynamicCast<MediaDataEx> (pSrc);
    // fill start codes info
    if (pSrcEx)
    {
        /* DEBUG: need to develop */
    }

    return UMC_OK;

} // Status AVSSampleBuffer::LockOutputBuffer(MediaData *pSrc)

Status AVSSampleBuffer::UnLockOutputBuffer(MediaData *pSrc)
{
    // check error(s)
    if (NULL == pSrc)
        return UMC_ERR_NULL_PTR;
    if (NULL == m_pbBuffer)
        return UMC_ERR_NOT_INITIALIZED;

    // need to update start codes array
    if (pSrc->GetDataSize())
    {
        if (pSrc->GetBufferSize() != pSrc->GetDataSize())
        {
            /* DEBUG: need to update start codes */
        }
    }
    else
    {
        m_latestSample.MoveDataPointer((Ipp32s) m_latestSample.GetDataSize());
    }

    // unlock the main buffer
    return m_mainBuffer.UnLockOutputBuffer(&m_latestSample);

} // Status AVSSampleBuffer::UnLockOutputBuffer(MediaData *pSrc)

void AVSSampleBuffer::CollectStartCodes(size_t nSourceSize)
{
    Ipp8u *pbSource = m_pbSampleEnd;

    // update end of the being collected sample
    m_pbSampleEnd += nSourceSize;

    // set up pointer to the beginning of the current frame
    {
        size_t iRollBack = IPP_MIN(pbSource - m_pbSample, 3);

        // roll back the pointer a little bit
        // to handle splitted start codes
        pbSource -= iRollBack;
    }

    // check amount of data
    if (4 > (m_pbSampleEnd - pbSource))
        return;

    // skip first 3 bytes
    pbSource += 3;

    // try to find all start codes
    do
    {
        // we can skip 3 bytes when
        // the previous byte is greater than 1
        if (1 < pbSource[-1])
        {
            pbSource += 3;
        }
        // we can skip 2 bytes when
        // two precending the previous byte byte is not 0
        else if (0 < pbSource[-2])
        {
            pbSource += 2;
        }
        // in other case we have to analyze stream byte by byte
        else
        {
            if (1 == pbSource[-1])
            {
                if (0 == pbSource[-3])
                {
                    // fill start code's info
                    m_startCodes[m_iNumberOfStartCodes].m_startCode = pbSource[0];
                    m_startCodes[m_iNumberOfStartCodes].m_startCodeOffset = (Ipp32u) (pbSource - 3 - m_pbSample);

                    // increment number of start codes
                    m_iNumberOfStartCodes += 1;

                    pbSource += 3;
                }
                else
                {
                    pbSource += 2;
                }
            }

            pbSource += 1;
        }

    } while (pbSource < m_pbSampleEnd);

    // we have read garbage
    if (0 == m_iNumberOfStartCodes)
    {
        size_t iRollBack = IPP_MIN(m_pbSample - m_pbBuffer, 3);

        m_pbSample = m_pbSampleEnd - iRollBack;
    }

} // void AVSSampleBuffer::CollectStartCodes(size_t nSourceSize)

Status AVSSampleBuffer::MoveCollectedSample(Status streamStatus)
{
    Ipp32s iNALUnits;
    Status umcRes;

    // sometimes we need to set 'EOB' having zero data
    if ((UMC_OK != streamStatus) &&
        (0 == m_iNumberOfStartCodes))
    {
        m_mainBuffer.UnLockInputBuffer(NULL, UMC_ERR_END_OF_STREAM);
        return UMC_OK;
    }

    // move collected samples to the main buffer
    iNALUnits = GetNALUnitsNumber();
    while (iNALUnits)
    {
        size_t nSampleSize, nEncodedDataSize;
        MediaData mediaDst;
        Ipp32s i;

        // calculate the sample's size
        nEncodedDataSize = m_startCodes[iNALUnits].m_startCodeOffset -
                           m_startCodes[0].m_startCodeOffset;
        nSampleSize = align_value<size_t> (nEncodedDataSize, sizeof(Ipp32u)) +
                      sizeof(Ipp32u) * (iNALUnits * 2 + 2);

        // get the destination
        umcRes = m_mainBuffer.LockInputBuffer(&mediaDst, nSampleSize);
        if (UMC_OK != umcRes)
            return umcRes;

        // copy sample & start codes' info
        {
            Ipp8u *pbDst = (Ipp8u *) mediaDst.GetBufferPointer();
            Ipp32u *pCodes;

            // copy sample
            memcpy(pbDst, m_pbSample + m_startCodes[0].m_startCodeOffset, nEncodedDataSize);
            pbDst += nEncodedDataSize;

            // align destination pointer
            while ((pbDst - (Ipp8u *) 0) & 3)
            {
                *pbDst = 0;
                pbDst += 1;
            }

            // copy start codes
            pCodes = (Ipp32u *) pbDst;
            for (i = 0; i < iNALUnits; i += 1)
            {
                pCodes[i] = m_startCodes[i].m_startCode;
                pCodes[i + iNALUnits] = m_startCodes[i].m_startCodeOffset - m_startCodes[0].m_startCodeOffset;
            }
            // set the number of start codes
            pCodes[iNALUnits * 2] = iNALUnits;
            // set the real data size
            pCodes[iNALUnits * 2 + 1] = (Ipp32u) nEncodedDataSize;
        }

        // finalize the sample
        mediaDst.SetDataSize(nSampleSize);
        m_mainBuffer.UnLockInputBuffer(&mediaDst, streamStatus);

        // update array of start codes
        SkipStartCodes(iNALUnits);

        // try to find the next sample
        iNALUnits = GetNALUnitsNumber();
    }

    // update the temporal buffer, move collected data to the beginning
    if (m_pbSample != m_pbBuffer)
    {
        Ipp8u *pbSrc, *pbDst;
        size_t nAvailable, nMaxChunk;

        nAvailable = m_pbSampleEnd - m_pbSample;
        nMaxChunk = m_pbSample - m_pbBuffer;
        pbDst = m_pbBuffer;
        pbSrc = m_pbSample;

        while (nAvailable)
        {
            size_t nToCopy = IPP_MIN(nMaxChunk, nAvailable);

            memcpy(pbDst, pbSrc, nToCopy);
            pbDst += nToCopy;
            pbSrc += nToCopy;
            nAvailable -= nToCopy;
        }

        m_pbSampleEnd = m_pbBuffer + (m_pbSampleEnd - m_pbSample);
        m_pbSample = m_pbBuffer;
    }

    return UMC_OK;

} // Status AVSSampleBuffer::MoveCollectedSample(Status streamStatus)

void AVSSampleBuffer::SkipStartCodes(Ipp32s numCodes)
{
    Ipp32s i;
    size_t nAdvanceSize;

    // check the following code
    if (VIDEO_SEQUENCE_END_CODE == m_startCodes[numCodes].m_startCode)
    {
        // skip filled buffer
        nAdvanceSize = m_startCodes[numCodes].m_startCodeOffset + sizeof(Ipp32u);
        numCodes += 1;
    }
    else
    {
        nAdvanceSize = m_startCodes[numCodes].m_startCodeOffset;
    }

    // skip filled buffer
    m_pbSample += nAdvanceSize;

    // skip codes info
    for (i = 0; i < (m_iNumberOfStartCodes - numCodes); i += 1)
    {
        m_startCodes[i].m_startCode = m_startCodes[i + numCodes].m_startCode;
        m_startCodes[i].m_startCodeOffset = (Ipp32u) (m_startCodes[i + numCodes].m_startCodeOffset - nAdvanceSize);
    }
    m_iNumberOfStartCodes -= numCodes;

} // void AVSSampleBuffer::SkipStartCodes(Ipp32s numCodes)

Ipp32s AVSSampleBuffer::GetNALUnitsNumber(void)
{
    Ipp32s i;
    bool bFoundPictureHeader = false;
    bool bFoundSequenceHeader = false;
    bool bFoundSlice = false;
    bool bPictureDone = false;

    if (0 >= m_iNumberOfStartCodes)
        return 0;

    // find the first slice start code or second other header
    for (i = 0; i < m_iNumberOfStartCodes; i += 1)
    {
        switch (m_startCodes[i].m_startCode)
        {
        case VIDEO_SEQUENCE_START_CODE:
            // the picture is complete, when it is not the first NAL unit
            if (i)
                bPictureDone = true;
            else
            {
                bFoundSequenceHeader = true;
                m_bFoundSequenceHeader = true;
            }
            break;

        case VIDEO_SEQUENCE_END_CODE:
            bPictureDone = true;
            break;

        case USER_DATA_START_CODE:
            // picture data before user data is not allowed
            if (bFoundSlice)
                bPictureDone = true;
            break;

        case I_PICTURE_START_CODE:
            // two header per frame is not allowed
            if ((bFoundSlice) || (bFoundPictureHeader))
                bPictureDone = true;
            else
                bFoundPictureHeader = true;
            break;

        case EXTENSION_START_CODE:
            // picture data before extension data is not allowed
            if (bFoundSlice)
                bPictureDone = true;
            break;

        case BP_PICTURE_START_CODE:
            if ((bFoundSlice) || (bFoundPictureHeader))
                bPictureDone = true;
            else
                bFoundPictureHeader = true;
            break;

        case VIDEO_EDIT_CODE:
            // video edit code shall be the beginning
            if ((bFoundSlice) || (bFoundPictureHeader))
                bPictureDone = true;
            break;

        // slices data
        default:
            // we should provide a sequence header as a single frame
            if ((bFoundSequenceHeader) && (false == bFoundPictureHeader))
                bPictureDone = true;
            else if (LAST_SLICE_START_CODE >= m_startCodes[i].m_startCode)
                bFoundSlice = true;
            break;
        }

        // is the picture composed
        if (bPictureDone)
        {
            // frame is incomplete.
            // there are only two conditions:
            // we found a sequence header or
            // we found a picture header and a sequence header had been found before
            if ((bFoundSequenceHeader) ||
                (bFoundPictureHeader && m_bFoundSequenceHeader))
                return i;
            else
            {
                SkipStartCodes(i);

                // try to compose another frame
                i = -1;
                bFoundPictureHeader = false;
                bFoundSequenceHeader = false;
                bFoundSlice = false;
                bPictureDone = false;
            }
        }
    }

    return 0;

} // Ipp32s AVSSampleBuffer::GetNALUnitsNumber(void)

} // namespace UMC

#endif // #if defined(UMC_ENABLE_AVS_SPLITTER)
